import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

public class Agent {
	private HashMap<State, HashMap<Direction, Float>> qValues;
	private HashMap<State, HashMap<Direction, Float>> dirtyQValues;
	private HashMap<State, Boolean> dirtyStates;
	private HashMap<Direction, Float> sortedDirections;
	private float alfa;
	private float gamma;
	private EpsilonPolicy policy;
	private State start;
	private State goal;
	private State current;
	private GridWorld world;
	private final int STANDARD_REWARD;
	private final int GOAL_REWARD;
	private final int COLLISION_REWARD;
	public BufferedWriter out;

	public Agent(float alfa, float gamma, float epsilon, GridWorld world,
			Coordinate start, Coordinate goal, int standardReward,
			int goalReward, int collisionReward, String name) {
		this.alfa = alfa;
		this.gamma = gamma;
		this.policy = new EpsilonPolicy(epsilon);
		this.world = world;
		this.start = this.getWorld().getState(start);
		this.current = this.getWorld().getState(start);
		this.goal = this.getWorld().getState(goal);
		this.STANDARD_REWARD = standardReward;
		this.GOAL_REWARD = goalReward;
		this.COLLISION_REWARD = collisionReward;
		this.initQValues();
		FileWriter fstream;
		try {
			fstream = new FileWriter(name);
			out = new BufferedWriter(fstream);
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	private void initQValues() {
		this.qValues = new HashMap<State, HashMap<Direction, Float>>();
		this.dirtyQValues = new HashMap<State, HashMap<Direction, Float>>();
		this.dirtyStates = new HashMap<State, Boolean>();
		State[][] states = this.getWorld().getStates();
		for (int i = 0; i < states.length; i++) {
			for (int j = 0; j < states[0].length; j++) {
				this.dirtyStates.put(states[i][j], false);
				HashMap<Direction, Float> values = new HashMap<Direction, Float>();
				for (Direction direction : Direction.values()) {
					values.put(direction, 0.0f);
				}
				this.qValues.put(states[i][j], values);
			}
		}
	}

	private GridWorld getWorld() {
		return this.world;
	}

	private float calculateNewQValue(State newState, Direction direction,
			int reward) {
		float result = this.getQValue(this.getCurrentState(), direction);
		float maxQValue = this.getMaxQValue(newState);
		maxQValue *= this.getGamma();
		maxQValue += reward;
		maxQValue -= this.getQValue(this.getCurrentState(), direction);
		maxQValue *= this.getAlfa();
		result += maxQValue;
		return result;
	}

	private float getMaxQValue(State newState) {
		float result = this.getQValue(newState, Direction.NORTH);
		for (Direction direction : Direction.values()) {
			if (this.getQValue(newState, direction) > result) {
				result = this.getQValue(newState, direction);
			}
		}
		return result;
	}

	private float getGamma() {
		return this.gamma;
	}

	private float getAlfa() {
		return this.alfa;
	}

	protected void takeSingleLearningStep() {
		// Kies action volgens epsilon-greedy policy
		Direction direction = policy.getPreferredDirection(this.getQValues(this
				.getCurrentState()));

		State newState = this.getNewState(direction);

		if (newState == null) {
			newState = this.getCurrentState();
		}

		// Neem die action en sla reward en nieuwe state op
		int reward;
		if (newState.equals(getGoalState()))
			reward = this.GOAL_REWARD;
		else
			reward = this.STANDARD_REWARD;

		// Bereken nieuwe Q-value
		float newQValue = this.calculateNewQValue(newState, direction, reward);
		this.setQValue(this.getCurrentState(), direction, newQValue);

		// verwissel states
		this.setCurrentState(newState);
	}

	private State getCurrentState() {
		return this.current;
	}

	public boolean isInGoal() {
		return this.getCurrentState().equals(this.getGoalState());
	}

	private State getGoalState() {
		return this.goal;
	}

	private State getStartState() {
		return this.start;
	}

	public void reset() {
		this.setCurrentState(this.getStartState());
	}

	public String bestDirections() {
		String result = "***Directions*** \n";
		for (int i = this.getWorld().getStates()[0].length - 1; i >= 0; i--) {
			for (int j = 0; j < this.getWorld().getStates().length; j++) {
				State state = this.getWorld().getStates()[j][i];
				if (state == null) {
					result = result.concat("(NULL)");
				} else {
					result = result.concat(this.bestDirection(state));
				}
			}
			result = result.concat("\n");
		}
		return result;
	}

	private void setCurrentState(State state) {
		this.current = state;
	}

	protected void setQValue(State state, Direction direction, float value) {
		this.qValues.get(state).put(direction, value);
	}

	protected void setDirtyQValue(State state, Direction direction, float value) {
		this.dirtyQValues.get(state).put(direction, value);
	}

	protected float getQValue(State state, Direction direction) {
		return this.qValues.get(state).get(direction);
	}

	protected float getDirtyQValue(State state, Direction direction) {
		return this.dirtyQValues.get(state).get(direction);
	}

	protected HashMap<Direction, Float> getQValues(State state) {
		return this.qValues.get(state);
	}

	public String bestDirection(State state) {
		String result = "(";
		Direction bestDirection = Direction.NORTH;
		for (Direction direction : Direction.values()) {
			if (this.qValues.get(state).get(direction) > this.qValues
					.get(state).get(bestDirection)) {
				bestDirection = direction;
			}
		}
		result = result.concat(bestDirection + ")");
		return result;
	}

	public Direction getBestDirection(State state, boolean dirty, int index) {
		Direction result = Direction.NORTH;
		if (dirty) {
			return this.getDirection(index);
		} else {
			for (Direction direction : Direction.values()) {
				if (this.getQValue(getCurrentState(), direction) > this
						.getQValue(getCurrentState(), result)) {
					result = direction;
				}
			}
			return result;
		}
	}

	public Direction getBestDirection(State state, boolean dirty) {
		Direction result = Direction.NORTH;
		if (dirty) {
			for (Direction direction : Direction.values()) {
				if (this.getDirtyQValue(getCurrentState(), direction) > this
						.getDirtyQValue(getCurrentState(), result)) {
					result = direction;
				}
			}
			return result;
		} else {
			for (Direction direction : Direction.values()) {
				if (this.getQValue(getCurrentState(), direction) > this
						.getQValue(getCurrentState(), result)) {
					result = direction;
				}
			}
			return result;
		}
	}

	// public ArrayList<String> getPathToGoal() {
	// State currentState = this.getStartState();
	//
	// ArrayList<String> path = new ArrayList<String>();
	// String temp;
	//
	// while (!currentState.equals(this.getGoalState())) {
	// temp = this.bestDirection(currentState);
	// path.add(temp);
	// currentState = this.getNextState(currentState, temp);
	// }
	// return path;

	// State newState;
	// Direction direction;
	//
	// if(this.dirtyStates.get(this.getCurrentState())){
	// direction = this.getBestDirection(getCurrentState(), false);
	// newState = this.getNewState(direction);
	// if(newState.containsAgent()){
	// direction = this.getBestDirection(getCurrentState(), true);
	// }
	// }else{
	// direction = this.getBestDirection(getCurrentState(), false);
	// }
	// newState = this.getNewState(direction);
	// int index = 0;
	// while(newState.containsAgent()){
	// direction = this.getBestDirection(this.getCurrentState(), true, index);
	// index++;
	// }
	// if(newState.containsAgent()){
	// }
	// }

	private State getNextState(State current, String direction) {
		int xCoord = current.getX();
		int yCoord = current.getY();

		if (direction.equals("(NORTH)"))
			return this.getWorld().getStates()[xCoord][yCoord + 1];
		else if (direction.equals("(EAST)"))
			return this.getWorld().getStates()[xCoord + 1][yCoord];
		else if (direction.equals("(SOUTH)"))
			return this.getWorld().getStates()[xCoord][yCoord - 1];
		else if (direction.equals("(WEST)"))
			return this.getWorld().getStates()[xCoord - 1][yCoord];

		return null;
	}

	public void takeCoordinatedLearningStep() {
		Direction direction;
		State newState;
		// is nieuwe state dirty?
		if (this.isCurrentStateDirty()) {
			// ja, dan joint tabel gebruiken en updaten
			direction = this.getBestDirection(this.getCurrentState(), false);
			newState = this.getNewState(direction);
			if (newState == null) {
				newState = this.getCurrentState();
			}
			// TODO: kan beter door bij te houden waar andere agent vorige keer
			// stond??
			if (newState != null && newState.containsAgent()) {
				direction = this.getBestDirection(getCurrentState(), true);
			}
		} else {
			// nee, dan gewone tavel gebruiken en niet updaten
			direction = this.getBestDirection(this.getCurrentState(), false);
		}

		newState = getNewState(direction);
		if (newState == null) {
			newState = this.getCurrentState();
		}

		// collision?
		int reward;
		if (newState != null && newState.containsAgent()) {
			// ja, dan indien nodig markeren
			reward = this.COLLISION_REWARD;
			boolean firstTimeDirty = this.markStateAsDirty(this
					.getCurrentState());
			if (firstTimeDirty) {
				// init dirty q value
				this.initDirtyStateQValues(this.getCurrentState());
			}
			// nee, gewoon voordoen
		} else if (this.isInGoal()) {
			reward = this.GOAL_REWARD;
		} else {
			reward = this.STANDARD_REWARD;
		}

		// indien nodig dirty joint tabel updaten.
		if (this.isCurrentStateDirty()) {
			float newQValue = this.calculateDirtyQValue(newState, direction,
					reward);
			this.setDirtyQValue(this.getCurrentState(), direction, newQValue);
		}

		// agent van plaats veranderen als er geen collision zou plaatsvinden
		if (!(reward == this.COLLISION_REWARD)) {
			this.getCurrentState().removeAgent();
			this.setCurrentState(newState);
			this.getCurrentState().addAgent();
		}
	}

	private float calculateDirtyQValue(State newState, Direction direction,
			int reward) {
		float result = this.getDirtyQValue(this.getCurrentState(), direction);
		float maxQValue = this.getMaxQValue(newState);
		maxQValue *= this.getGamma();
		maxQValue += reward;
		maxQValue -= this.getDirtyQValue(this.getCurrentState(), direction);
		maxQValue *= this.getAlfa();
		result += maxQValue;
		return result;
	}

	private void initDirtyStateQValues(State state) {
		HashMap<Direction, Float> values = new HashMap<Direction, Float>();
		for (Direction direction : Direction.values()) {
			values.put(direction, 0.0f);
		}
		this.dirtyQValues.put(state, values);
	}

	private boolean isCurrentStateDirty() {
		return this.dirtyStates.get(this.getCurrentState());
	}

	private boolean markStateAsDirty(State currentState) {
		if (this.dirtyStates.get(currentState)) {
			return false;
		} else {
			this.dirtyStates.put(currentState, true);
			return true;
		}
	}

	private State getNewState(Direction direction) {
		int x = this.getCurrentState().getX() + direction.getXDisplacement();
		int y = this.getCurrentState().getY() + direction.getYDisplacement();
		Coordinate coordinate = new Coordinate(x, y);
		return this.getWorld().getState(coordinate);
	}

	public void enterStartState() {
		if (!this.getCurrentState().equals(this.getStartState())) {
			throw new NullPointerException(
					"Agent must be in startposition to begin.");
		}
		this.getCurrentState().addAgent();
		this.getGoalState().removeAgent();
	}

	public void takeNormalStep() {
		State newState;
		Direction direction;

		if (this.dirtyStates.get(this.getCurrentState())) {
			direction = this.getBestDirection(getCurrentState(), false);
			newState = this.getNewState(direction);
			if (newState != null && newState.containsAgent()) {
				direction = this.getBestDirection(getCurrentState(), true);
			}
		} else {
			direction = this.getBestDirection(getCurrentState(), false);
		}
		newState = this.getNewState(direction);
		if (newState != null && this.getNewState(direction).containsAgent()) {
			// teller++;
			int index = 1;
			while (this.getNewState(direction).containsAgent()) {
				this.sortDirections();
				try {
					direction = this.getBestDirection(getCurrentState(), true,
							index);
				} catch (Exception e) {
					break;
				}
				newState = this.getNewState(direction);
				index++;
			}
		}
		if (newState == null) {
			newState = this.getCurrentState();
		}
		this.write(direction.toString());
		this.getCurrentState().removeAgent();
		this.setCurrentState(newState);
		this.getCurrentState().addAgent();
	}

	public void sortDirections() {
		sortedDirections = new LinkedHashMap<Direction, Float>();
		List<Direction> mapKeys = new ArrayList<Direction>(this
				.getDirtyQValues(this.getCurrentState()).keySet());
		List<Float> mapValues = new ArrayList<Float>(this.getDirtyQValues(
				this.getCurrentState()).values());
		TreeSet<Float> sortedSet = new TreeSet<Float>(mapValues);
		Object[] sortedArray = sortedSet.toArray();
		int size = sortedArray.length;
		for (int i = 0; i < size; i++) {
			sortedDirections.put(
					mapKeys.get(mapValues.indexOf(sortedArray[i])),
					(Float) sortedArray[i]);
		}
	}

	private HashMap<Direction, Float> getDirtyQValues(State currentState) {
		return this.dirtyQValues.get(currentState);
	}

	public Direction getDirection(int index) {
		Set<Direction> ref = this.sortedDirections.keySet();
		Iterator<Direction> it = ref.iterator();
		for (int i = 0; i < index; i++) {
			it.next();
		}
		return it.next();
	}

	public static int teller = 0;

	public void write(String string) {
		try {
			out.write(string+"\n");
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
}
